package com.wywah.yunduo.security.supports.custom;

import com.wywah.yunduo.security.utils.SpringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.support.MessageSourceAccessor;
import org.springframework.security.authentication.AbstractAuthenticationToken;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;

import java.security.Principal;
import java.util.*;

public abstract class YunduoCustomAuthenticationProvider<T extends YunduoCustomAuthenticationToken> implements AuthenticationProvider {
    private static final String ERROR_URI = "https://datatracker.ietf.org/doc/html/rfc6749#section-4.1.2.1";

    private Logger logger = LoggerFactory.getLogger(getClass());

    private final OAuth2AuthorizationService authorizationService;

    private final OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator;

    private final AuthenticationManager authenticationManager;

    public YunduoCustomAuthenticationProvider(OAuth2AuthorizationService authorizationService, OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator, AuthenticationManager authenticationManager) {
        Assert.notNull(authorizationService, "authorizationService cannot be null");
        Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
        this.authorizationService = authorizationService;
        this.tokenGenerator = tokenGenerator;
        this.authenticationManager = authenticationManager;
    }

    public YunduoCustomAuthenticationProvider(HttpSecurity http){
        Assert.notNull(http, "http cannot be null");
        this.authorizationService = http.getSharedObject(OAuth2AuthorizationService.class);
        this.tokenGenerator = http.getSharedObject(OAuth2TokenGenerator.class);
        this.authenticationManager = http.getSharedObject(AuthenticationManager.class);
    }


    /**
     * 当前provider是否支持此令牌类型
     * @param authentication
     * @return
     */
    @Override
    public abstract boolean supports(Class<?> authentication);

    /**
     * 当前的请求客户端是否支持此模式
     * @param registeredClient
     */
    public abstract void verifyClient(RegisteredClient registeredClient);

    public abstract AbstractAuthenticationToken buildToken(Map<String, Object> reqParameters);

    @Override
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {

        T customAuthentication = (T) authentication;

        OAuth2ClientAuthenticationToken clientPrincipal = getAuthenticatedClientElseThrowInvalidClient(
                customAuthentication);

        RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();
        // 验证客户端
        verifyClient(registeredClient);

        //处理授权范围
        Set<String> authorizedScopes;
        if (!CollectionUtils.isEmpty(customAuthentication.getScopes())) {
            for (String requestedScope : customAuthentication.getScopes()) {
                if (!registeredClient.getScopes().contains(requestedScope)) {
                    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_SCOPE);
                }
            }
            authorizedScopes = new LinkedHashSet<>(customAuthentication.getScopes());
        }
        else {
            authorizedScopes = new LinkedHashSet<>();
        }
        Map<String, Object> reqParameters = customAuthentication.getAdditionalParameters();
        try {
            AbstractAuthenticationToken authenticationToken = buildToken(reqParameters);
            Authentication newAuthentication = authenticationManager
                    .authenticate(authenticationToken);

            // @formatter:off
            DefaultOAuth2TokenContext.Builder tokenContextBuilder = DefaultOAuth2TokenContext.builder()
                    .registeredClient(registeredClient)
                    .principal(newAuthentication)
                    .authorizationServerContext(AuthorizationServerContextHolder.getContext())
                    .authorizedScopes(authorizedScopes)
                    .authorizationGrantType(customAuthentication.getAuthorizationGrantType())
                    .authorizationGrant(customAuthentication)
                    .context(context -> context.put("reqParameters", reqParameters));
            // @formatter:on
            
            OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization
                    .withRegisteredClient(registeredClient)
                    .principalName(newAuthentication.getName())
                    .authorizationGrantType(customAuthentication.getAuthorizationGrantType())
                    // 0.4.0 新增的方法
                    .authorizedScopes(authorizedScopes);


            // ----- Access token -----
            OAuth2TokenContext tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
            OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
            if (generatedAccessToken == null) {
                OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                        "The token generator failed to generate the access token.", ERROR_URI);
                throw new OAuth2AuthenticationException(error);
            }

            OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
                    generatedAccessToken.getTokenValue(), generatedAccessToken.getIssuedAt(),
                    generatedAccessToken.getExpiresAt(), tokenContext.getAuthorizedScopes());
            if (generatedAccessToken instanceof ClaimAccessor) {
                authorizationBuilder.id(accessToken.getTokenValue())
                        .token(accessToken,
                                (metadata) -> metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME,
                                        ((ClaimAccessor) generatedAccessToken).getClaims()))
                        // 0.4.0 新增的方法
                        .authorizedScopes(authorizedScopes)
                        .attribute(Principal.class.getName(), newAuthentication)
                        .attribute(OAuth2ParameterNames.ACCESS_TOKEN, OAuth2ParameterNames.ACCESS_TOKEN);
            }
            else {
                authorizationBuilder.id(accessToken.getTokenValue()).accessToken(accessToken);
            }
            /*----- Refresh token -----*/
            OAuth2RefreshToken refreshToken = null;
            // 客户端支持 Refresh token 并且不是公共客户端
            if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
                    !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

                tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
                OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
                if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
                    OAuth2Error error = new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR,
                            "The token generator failed to generate the refresh token.", ERROR_URI);
                    throw new OAuth2AuthenticationException(error);
                }
                refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
                authorizationBuilder.refreshToken(refreshToken).attribute(OAuth2ParameterNames.REFRESH_TOKEN, OAuth2ParameterNames.REFRESH_TOKEN);
            }

            OAuth2Authorization authorization = authorizationBuilder.build();
            // 保存授权信息
            this.authorizationService.save(authorization);

            return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, accessToken,
                    refreshToken, Objects.requireNonNull(authorization.getAccessToken().getClaims()));
        }catch (Exception e){
            logger.error("授权异常", e);
            throw new OAuth2AuthenticationException(new OAuth2Error("authentication error",
                    "授权错误",
                    ""));
        }
    }

    private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(
            Authentication authentication) {

        OAuth2ClientAuthenticationToken clientPrincipal = null;

        if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
            clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
        }

        if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
            return clientPrincipal;
        }

        throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
    }
}
